"""rnn in pytorch
"""

import torch
import torch.nn as nn


class RNNModel(nn.Module):

    def __init__(self, input_size, output_size, hidden_number, layer_number, device=None):
        super().__init__()
        self.hidden_number = hidden_number
        self.layer_number = layer_number
        self.device = device
        self.rnn = nn.RNN(input_size=input_size, hidden_size=hidden_number, num_layers=layer_number, batch_first=True)
        self.linear = nn.Linear(hidden_number, output_size)
        self.relu = nn.ReLU()

    def forward(self, x):
        batch_size = x.shape[0]
        state = self.begin_state(batch_size=batch_size)
        out, state = self.rnn(x, state)
        output = self.linear(self.relu(out))

        return output, state

    def begin_state(self, batch_size):
        if self.device:
            return torch.zeros(self.layer_number, batch_size, self.hidden_number).to(self.device)
        else:
            return torch.zeros(self.layer_number, batch_size, self.hidden_number)
